import os
import subprocess
from tqdm import tqdm
import time
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from configure import *


class ida_disassembly_tool:
    def __init__(self, script_location, samplefile, savePath, worknum=1) -> None:
        self.script = script_location
        self.samplefile = samplefile
        self.workNum = worknum
        self.idaPath = "D:\software\IDA_Pro_7.7\ida64.exe"
        self.scriptPath = script_location
        self.space = "      "
        self.savePath = savePath
        self.totalNum = 0
        self.sampleList = {}
        self.finishNum = 0
        self.finishList = {}
        self.unfinishedNum = 0
        self.unfinishedList = {}
        self.find_sample_to_disassembly()

    def analysisSample(self, samplePath, loc):
        """  
            运行如下命令
            D:\software\IDA\Hex-RaysIDA\ida64.exe -B  -SD:\code\pycode\ida\1.py D:\code\pycode\ida\1.exe
        """
        a = time.time()
        cmd = "{0} -L{1} -A -S{2} {3}".format(self.idaPath,
                                              ida_log, self.scriptPath, samplePath)
        print(cmd)
        # cmd=self.idaPath+" -L{} -A -S"+self.scriptPath+" "+samplePath
        subprocess.call(cmd)
        b = time.time()
        write_to_logs_file(
            disassemblylogPath, "{}{}处理完成****{}耗时: {}s\n".format(self.space, samplePath, self.space, b-a))
        time.sleep(2)

    # 找到需要处理的样本名，支持断点
    def find_sample_to_disassembly(self):
        for file in os.listdir(self.samplefile):
            data = file.split(".")  # 找到样本文件
            if (len(data) == 1):
                self.totalNum += 1
                self.sampleList[data[0]] = True
        for file in os.listdir(self.savePath):
            data = file.split(".")
            if (data[1] == "bak"):
                self.finishNum += 1
                self.finishList[data[0]] = True
        for name in self.sampleList.keys():
            if name not in self.finishList:
                self.unfinishedNum += 1
                self.unfinishedList[name] = self.samplefile+"//"+name
        write_to_logs_file(disassemblylogPath, "样本总计{}个，已处理{}个,还剩{}个\n".format(
            self.totalNum, self.finishNum, self.unfinishedNum))

    def run(self):
        thread_pool = ThreadPoolExecutor(max_workers=self.workNum)
        loc = 0
        for name in self.unfinishedList.keys():
            loc += 1
            samplePath = self.unfinishedList[name]
            thread_pool.submit(self.analysisSample, samplePath, loc)
        thread_pool.shutdown(wait=True)


def write_to_logs_file(disassemblylogPath, content):
    with open(disassemblylogPath, 'a+', encoding="utf-8") as file:
        file.write(str(datetime.now())+str(content)+"\n")


if __name__ == '__main__':
    idadis = ida_disassembly_tool(r"E:\code\ida\disassembly_with_ida\ida_script.py",
                                  r"E:\code\share\functionSImData\data\all_samples",
                                  r"E:\code\ida\disassembly_with_ida\data\finalData")
    idadis.run()
